package org.apache.lorry.zookeeper.cluster;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

@Slf4j
public class Client {
    private String ipPort;

    private ZooKeeper zk;

    public static final int SESSION_TIMEOUT = 10000;

    private volatile boolean connected = false;

    private volatile boolean expired = false;

    protected ConcurrentHashMap<String, Object> ctxMap = new ConcurrentHashMap<String, Object>();

    public Client(String ipPort) {
        this.ipPort = ipPort;
    }

    public void start() throws IOException, InterruptedException {
        CountDownLatch latch = new CountDownLatch(1);
        zk = new ZooKeeper(ipPort, SESSION_TIMEOUT, (event) -> {
            switch (event.getState()) {
                case SyncConnected:
                    connected = true;
                    latch.countDown();
                    break;
                case Disconnected:
                    connected = false;
                    log.error("Disconnected");
                    break;
                case Expired:
                    connected = false;
                    expired = true;
                    System.out.println("Expired");
                    break;
                default:
                    break;
            }
        });
        latch.await(SESSION_TIMEOUT, TimeUnit.MICROSECONDS);
    }

    public void submitTask(Task task) {
        try {
            zk.create(ZooConstant.TASK_PATH + "/" + task.getName(), JSON.toJSONString(task).getBytes(),
                    ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, createTaskCallback, task);
        } catch (Exception e) {
            log.error("submitTask error", e);
        }
    }


    AsyncCallback.StringCallback createTaskCallback = (rc, path, ctx, name) -> {
        Task myTask = (Task) ctx;
        switch (KeeperException.Code.get(rc)) {
            case CONNECTIONLOSS:
                submitTask(myTask);
                break;
            case OK:
                log.info("creat task : {} OK", myTask.getName());
                watchStatus(myTask);
                break;
            default:
                log.error("Something went wrong: {}", KeeperException.create(KeeperException.Code.get(rc), path));
                break;
        }
    };

    private void watchStatus(Task myTask) {
        ctxMap.put(myTask.getName(), myTask);
        zk.exists(ZooConstant.STATUS_PATH + "/" + myTask.getName(), statusWatcher, statCallback, myTask);
    }

    Watcher statusWatcher = event -> {
        if (event.getType() == Watcher.Event.EventType.NodeCreated) {
            log.info("Created task: {}", event.getPath());
            String task = event.getPath().replace(ZooConstant.STATUS_PATH + "/", "");
            zk.getData(event.getPath(), false, this.getDataCallback, ctxMap.get(task));
        }
    };


    AsyncCallback.StatCallback statCallback = new AsyncCallback.StatCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    watchStatus((Task) ctx);
                    break;
                case OK:
                    if (stat != null) {
                        zk.getData(path, false, getDataCallback, ctx);
                        log.info("Status node is there: {}", path);
                    }
                    break;
                case NONODE:
                    break;
                default:
                    log.error("Something went wrong when " +
                            "checking if the status node exists: " +
                            KeeperException.create(KeeperException.Code.get(rc), path));
                    break;
            }
        }
    };



    AsyncCallback.DataCallback getDataCallback = (rc, path, ctx, data, stat) -> {
        Task task = (Task) ctx;
        switch (KeeperException.Code.get(rc)) {
            case CONNECTIONLOSS:
                zk.getData(path, false, this.getDataCallback, ctxMap.get(path));
                break;
            case OK:
                String str = new String(data);
                log.info("Task {} data: {}", task.getName(), str);
                assert str.equals("done");
                task.done();
                //zk.delete("/tasks/" + path.replace("/status/", ""), -1, taskDeleteCallback, null);
                zk.delete(path, -1, this.taskDeleteCallback, null);
                ctxMap.remove(task.getName());
                break;
            case NONODE:
                log.info("Status node is gone!");
                break;
            default:
                log.error("Something went wrong here, " +
                        KeeperException.create(KeeperException.Code.get(rc), path));
                break;
        }
    };

    AsyncCallback.VoidCallback taskDeleteCallback = (rc, path, ctx) -> {
        switch (KeeperException.Code.get(rc)) {
        case CONNECTIONLOSS:
            zk.delete(path, -1, this.taskDeleteCallback, null);
            break;
        case OK:
            log.info("Successfully deleted {}", path);
            break;
        default:
            log.error("Something went wrong here, " +
                    KeeperException.create(KeeperException.Code.get(rc), path));
        }
    };


}
